import os
import struct
import numpy as np


def synthetic_data(features, labels):
    for i in range(len(labels)):
        yield features[i]/255, labels[i]


def load_mnist(path, split='train', reshape=False):
    """
    reference:
    https://blog.csdn.net/justidle/article/details/103146658
    """
    labels_path = os.path.join(path, f'{split}-labels-idx1-ubyte')
    images_path = os.path.join(path, f'{split}-images-idx3-ubyte')
    with open(labels_path, 'rb') as lb_path:
        magic, n = struct.unpack('>II', lb_path.read(8))
        labels = np.fromfile(lb_path, dtype=np.uint8)

    with open(images_path, 'rb') as img_path:
        magic, num, rows, cols = struct.unpack('>IIII', img_path.read(16))
        if reshape:
            images = np.fromfile(img_path, dtype=np.uint8).reshape(len(labels), 28, 28, 1)
        else:
            images = np.fromfile(img_path, dtype=np.uint8).reshape(len(labels), 28 * 28)
    return images, labels
